open Ast
open Sast
open Types
open Bigarray

(*
(* Make sure that:
 *
 *)
let rec Ast_expr = function


(* Make sure that:
 * The size of inputs are defined.
 *)
let rec Sast1_expr = function
 *)

exception ExceptionS2 of string;;


let str_of_ed_type = function
   Sast.Float(_) -> "Sast.Float"
 | Sast.Binop(_,_,_) -> "Sast.Binop"
 | Sast.Unop(op,_) -> 
     ( match op with
           Sast.Sin -> "Sast.Sin"
         | Sast.Cos -> "Sast.Cos"
         | Sast.Exp -> "Sast.Exp"
         | Sast.Negate -> "Sast.Negate"
         | Sast.Abs -> "Sast.Abs"
     )
 | Sast.Intgr(_) -> "Sast.Intgr"
 | Sast.SizeOf(_,_) ->"Sast.SizeOf"
 | Sast.Id(_) ->"Sast.Id"
 | Sast.NeuronRef(_,_,_) ->"Sast.NeuronRef"
 | Sast.ExtNeuron(_,_,_,_,_) ->"Sast.ExtNeuron"
 | Sast.Param(_,_) ->"Sast.Param"
 | Sast.ActRef(_,_,_,_) ->"Sast.ActRef"
 | Sast.WghtRef(_,_,_) ->"Sast.WghtRef"
 | Sast.Conv(_) ->"Sast.Conv"
 | Sast.Filter(_,_,_) ->"Sast.Filter"
 | Sast.Span(_,_,_) ->"Sast.Span"
 | Sast.ForExpr(_,_) ->"Sast.ForExpr"
 | Sast.Index(_,_) ->"Sast.Index"

let sast2str (e,t) = str_of_ed_type e

(* Need to make sure:
 * there are no macros (for, sizeof)
 * all dimensions of modules, neurons, and outputs are defined
 * no modules(?)
 *)
let rec sast2_expr = function
   exprd, Types.Matrix([]) -> 
     print_endline (str_of_ed_type exprd);
     assert false
 | Sast.WghtRef(_,_,_),Types.MatrixUnknownSize -> () (* A weight reference has infinite size *)
 | exprd, Types.MatrixUnknownSize -> 
      raise( ExceptionS2 ( (str_of_ed_type exprd)^" has unresolved size."))
 | Sast.Unop(_,x1), t -> sast2_expr x1
 | Sast.Binop(x1,_,x2), t -> sast2_expr x1; sast2_expr x2
 | Sast.Conv((x1,t1),(x2,t2)), Types.Matrix(diml) ->
     assert( t1=t2 || ( t1=Types.Int && t2=Types.Matrix([1]) ) );
     sast2_expr (x1,t1); 
     sast2_expr (x2,t2);
 | Sast.Filter(x1,x2,_), Types.Matrix(diml) -> 
     sast2_expr x1; sast2_expr x2
 | Sast.Filter(_,_,_),Types.Int -> assert false
 | Sast.Index(x1,_), Types.Matrix(diml) -> 
     sast2_expr x1
 | Sast.WghtRef(_,_,_),Types.Int -> assert false
 | Sast.WghtRef(_,_,_),Types.Matrix(diml) -> ()
 | Sast.Conv(x1,x2), Types.Int ->  assert false
 | Sast.Intgr(_), Types.Int -> ()
 | Sast.Intgr(_), t -> raise( ExceptionS2 "Intgr has been associated with a non-int type.")
 | Sast.Float(_), Types.Matrix([1]) -> ()
 | Sast.Float(_), t -> raise( ExceptionS2 "Float has been associated with a non-float type (ie. single element matrix).")
 | Sast.SizeOf(exprd,_),Types.Int -> () (* more checks? *)
 | Sast.SizeOf(exprl,_),_ -> raise( ExceptionS2 "SizeOf has been associated with a non-int type.")
 | Sast.Id(_),_ -> ()
 | Sast.Param(pnum,diml), Types.Matrix(_) -> ()
 | Sast.Param(pnum,diml), _ -> raise( ExceptionS2 "Possible invalid Sast2 progrm")
 | Sast.ActRef(-1,_,_,_), _ -> raise( ExceptionS2 "An activation function reference is unresolved.")
 | Sast.ActRef(gsid,name,expr,fparamVals), Types.Matrix([1]) -> ()
 | Sast.ActRef(gsid,name,expr,fparamVals), Types.Matrix(_) -> 
     raise( ExceptionS2 "Activation function must be scalar." )
 | Sast.ActRef(gsid,name,expr,fparamVals), _ -> raise( ExceptionS2 "Type/size of ActRef is unresolved.")
 | Sast.ExtNeuron(-1,mname,_,nname,_), _ -> 
     raise( ExceptionS2 
              (Printf.sprintf "Unresolve reference to external neuron %s.%s"
                mname nname ))
 | Sast.ExtNeuron(_,_,_,_,_), Types.Matrix(_)-> ()
 | Sast.ExtNeuron(_,_,_,_,_), _ -> raise( ExceptionS2 "Type/size of ExtNeuron is unresolved.")
 | Sast.NeuronRef(-1,name,indl), _ -> 
     raise( ExceptionS2 
              (Printf.sprintf "Unresolve reference to local neuron '%s'." name ))
 | Sast.NeuronRef(gsid,name,indl), Types.Matrix(_) -> ()
 | n -> raise( ExceptionS2 ("Validation code not completed for "^(sast2str n)^"."))

let rec sast2_forexpr = function 
    Sast.ForExpr(x1,x2) -> ()
  | e -> raise( ExceptionS2 ("Expected ForExpr but found "^(str_of_ed_type e)^"."))

let sast2_syn syn =
  sast2_expr syn.sexpr;
  List.iter sast2_forexpr syn.sfor

let sast2_param param = 
  match param with
      { pa_type = Matrix(diml); pa_data = ImageC3(data) } -> 
        let diml' = [ (Array3.dim1 data); (Array3.dim2 data); (Array3.dim3 data) ] 
        in
          if diml' <> diml then
            raise( ExceptionS2 ("PPM parameter's array has an inconsistent size."))
    | { pa_type = Matrix(diml); pa_data = FloatArray(data) } -> ()
    | { pa_type = Matrix(diml); pa_data = Unallocated } -> 
        raise( ExceptionS2 ("Parameter data unallocated."))
    | { pa_type = MatrixUnknownSize } -> raise( ExceptionS2 ("Parameter's size is unresolved."))
    | { pa_type = Int } -> raise( ExceptionS2 ("Parameter cannot be of type int."))


let sast2_neuron k n = match n with
    {t=Matrix(diml); v=v} ->
      let reqd_size = List.fold_left (fun a dim -> a * dim ) 1 diml in
        if Array.length v <> reqd_size then
          raise( ExceptionS2 
                   (Printf.sprintf 
                      "Neuron requires array of size %d, however has size of %d." 
                      reqd_size (Array.length v)))
  | _ -> raise( ExceptionS2 ("Neuron must be of type Matrix"))

let sast2 p =
   IntMap.iter (fun k x -> sast2_syn x) p.syn;
   IntMap.iter (fun k x -> sast2_param x) p.pparams;
   IntMap.iter (fun k x -> sast2_neuron k x) p.pneurons;
   (* make sure every neuron is defined by a synapse *)
   IntMap.iter (fun gsid n ->
                  if IntMap.mem gsid p.syn = false then
                    raise( ExceptionS2
                             (Printf.sprintf
                                "Neuron %d isn't defined by any synapses!" gsid)))
     p.pneurons
